#include <boost/graph/adjacency_list.hpp>
#include "trm/mincut.h"

namespace trm {

namespace mincut {

using namespace boost;

class segmenter {
public:
	typedef size_t capacity;
	typedef adjacency_list_traits < vecS, vecS, directedS >
		::edge_descriptor edge_descriptor;
	typedef adjacency_list < 
		vecS, //OutEdgeList
		vecS, //VertexList
		directedS, //Directed //bidirectionalS?
		property < vertex_color_t, default_color_type >, //VertexProperties
		property < edge_capacity_t, capacity,
			property < edge_residual_capacity_t, capacity,
				property < edge_reverse_t, edge_descriptor > > >, //EdgeProperties
		no_property, //GraphProperties
		vecS //EdgeList
	> Graph;
	typedef Graph::vertex_descriptor vertex_descriptor;
	/// vmap stores vertex descriptors for each pixel of a frame
	typedef img::image<vertex_descriptor> vmap_type;
public:
	segmenter(mswin::timing_tree * timing, bool be_silent);
	size_t num_vertices() const
		{ return boost::num_vertices(g); }
	size_t num_edges() const
		{ return boost::num_edges(g); }
	/// fill vmap for a frame, for which belonging of some pixels is known,
	/// for undetermined pixels new vertices in graph are created
	void create_vmap(	
		mask const & obligatory_foreground_pixels,
		mask const & obligatory_background_pixels,
		vmap_type & vmap);
	/// fill vmap for a frame where bolonging of each pixel is known
	/// (vmap is filled only with source_vertex and sink_vertex);
	/// no new vertices in graph appear
	void create_vmap(mask const & m, vmap_type & vmap);
	/// adds new edges in graph among similar pixels of two images
	void establish_links(
		/// center image is an image, from which the center of a window is fetched
		byte_image const & center_image,
		/// periphery image is an image, from which all window pixels are fetched
		byte_image const & periphery_image,
		/// vertices map for the center image
		vmap_type const & center_vmap,
		/// vertices map for the periphery image
		vmap_type const & periphery_vmap,
		size_t hw, size_t nn, float aspect
	);
	/// find max-flow, then mark each vertex as belonging to source, to sink or ambiguous
	void solve();
	/// call this method after solve to retrieve mask for a frame (specified by its vmap),
	/// pass white_color to retrive mask of surely foreground pixels, black_color --- background, gray_color ---ambiguous
	void get_frame_mask(vmap_type const & vmap, default_color_type col, mask & m) const;
	/// call this method after solve to retrieve mask for a frame (specified by its vmap),
	/// each existing value in m is multiplied on 2, than 1 is added for ambiguous pixels, and 2 for surely foreground pixels
	void inc_frame_gray_mask(vmap_type const & vmap, gray_mask & m) const;
	/// do what the name say; the capacities in different directions become distinct
	void double_capacity_of_saturated_edges();
private:
	void add_edge(vertex_descriptor u, vertex_descriptor v, capacity c);
	/// return true if a new edge was created 
	bool inc_edge(vertex_descriptor u, vertex_descriptor v, capacity c);
	void add_bidir_edge(vertex_descriptor u, vertex_descriptor v, capacity c);
	// return true if new edges were created 
	bool inc_bidir_edge(vertex_descriptor u, vertex_descriptor v, capacity c);
private:
	Graph g;
	/// pixels specified in obligatory_foreground_pixels shares one vertex, and so for obligatory_background_pixels
	vertex_descriptor source_vertex, sink_vertex;
	property_map < Graph, vertex_color_t >::type			color;
	property_map < Graph, edge_capacity_t >::type			cap;
	property_map < Graph, edge_residual_capacity_t >::type	residual_capacity;
	property_map < Graph, edge_reverse_t >::type			rev_edge;
	mswin::timing_tree * timing_tree;
	bool silent;
};

inline void segmenter::add_edge(vertex_descriptor u, vertex_descriptor v, capacity c)
{
	assert (c > 0);
	assert (u != v);

	edge_descriptor e1, e1rev;
	
	e1 = boost::add_edge(u, v, g).first;
	e1rev = boost::add_edge(v, u, g).first;
	cap[e1] = c;
	cap[e1rev] = 0; 
	rev_edge[e1] = e1rev;
	rev_edge[e1rev] = e1;
}

inline bool segmenter::inc_edge(vertex_descriptor u, vertex_descriptor v, capacity c)
{
	assert (c > 0);
	assert (u != v);

	bool exists;
	edge_descriptor e;
	tie(e, exists) = edge(u, v, g);
	if (!exists) {
		add_edge(u, v, c);
		return true;
	}

	assert (cap[e] > 0);
	cap[e] += c;

	return false;
}

inline void segmenter::add_bidir_edge(vertex_descriptor u, vertex_descriptor v, capacity c)
{
	assert (c > 0);
	assert (u != v);

	edge_descriptor e1, e2, e1rev, e2rev;
	
	e1 = boost::add_edge(u, v, g).first;
	e2rev = boost::add_edge(u, v, g).first;
	e2 = boost::add_edge(v, u, g).first;
	e1rev = boost::add_edge(v, u, g).first;
	cap[e1] = cap[e2] = c;
	cap[e1rev] = cap[e2rev] = 0; 
	rev_edge[e1] = e1rev;
	rev_edge[e1rev] = e1;
	rev_edge[e2] = e2rev;
	rev_edge[e2rev] = e2;
}

inline bool segmenter::inc_bidir_edge(vertex_descriptor u, vertex_descriptor v, capacity c)
{
	assert (c > 0);
	assert (u != v);

	bool e1_exists, e2_exists;
	edge_descriptor e1, e2;

	tie(e1, e1_exists) = edge(u, v, g);
	if (!e1_exists) {
		add_bidir_edge(u, v, c);
		return true;
	}

	tie(e2, e2_exists) = edge(v, u, g);
	assert (e2_exists);
	assert (cap[e1] > 0);
	assert (cap[e1] == cap[e2]);
	cap[e1] += c;
	cap[e2] += c;

	return false;
}

} //mincut

} //param
